package org.mx.jwt.service.impl;

import java.io.FileInputStream;
import java.security.KeyFactory;
import java.security.KeyStore;
import java.security.interfaces.RSAPrivateCrtKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.RSAPublicKeySpec;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;

import org.mx.StringUtils;
import org.mx.TypeUtils;
import org.mx.error.UserInterfaceSystemErrorException;
import org.mx.jwt.error.UserInterfaceJwtErrorException;
import org.mx.jwt.service.JwtService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;

import com.auth0.jwt.JWT;
import com.auth0.jwt.JWTCreator;
import com.auth0.jwt.JWTVerifier;
import com.auth0.jwt.algorithms.Algorithm;
import com.auth0.jwt.interfaces.Claim;
import com.auth0.jwt.interfaces.DecodedJWT;

/**
 * 描述： JWT（JSON Web Tokens）服务实现类
 *
 * @author john peng
 * Date time 2018/8/19 下午7:21
 */
public class JwtServiceImpl implements JwtService {
	private static final Logger logger = LoggerFactory.getLogger(JwtServiceImpl.class);

	private Algorithm algorithm;
	private JWTVerifier verifier;

	@Value("${jwtauth.algorithm:}")
	private String algorithmStr;
	@Value("${jwtauth.issuer:mx institute}")
	private String issuer;
	@Value("${jwtauth.subject:Account authenticate}")
	private String subject;

	@Value("${jwtauth.hs.secret:secret word}")
	private String hsSecret;

	@Value("${jwtauth.rsa.keystore.path:}")
	private String rsaKeystorePath;
	@Value("${jwtauth.rsa.keystore.type:JKS}")
	private String rsaKeystoreType;
	@Value("${jwtauth.rsa.keystore.password1:edmund!@#123}")
	private String rsaPassword1;
	@Value("${jwtauth.rsa.keystore.password2:ds110119}")
	private String rsaPassword2;
	@Value("${jwtauth.rsa.keystore.alias:jetty}")
	private String rsaAlias;

	/**
	 * 根据指定的算法名称构造对应的算法对象
	 *
	 * @param algorithm 算法名称
	 * @return 算法对象
	 */
	private Algorithm getAlgorithm(String algorithm) {
		if (algorithm.startsWith("HS")) {
			return getHsAlgorithm(algorithm);
		} else if (algorithm.startsWith("RSA")) {
			return getRsaAlgorithm(algorithm);
		} else {
			if (logger.isErrorEnabled()) {
				logger.error(String.format("Unsupported algorithm: %s.", algorithm));
			}
			throw new UserInterfaceJwtErrorException(UserInterfaceJwtErrorException.JwtErrors.JWT_INITIALIZE_FAIL);
		}
	}

	/**
	 * 根据指定的HS（HMAC）算法名称，构造对应的RSA算法对象
	 *
	 * @param algorithm 算法名称
	 * @return HS算法对象
	 */
	private Algorithm getHsAlgorithm(String algorithm) {
		try {
			switch (algorithm) {
			case "HS512":
				return Algorithm.HMAC512(hsSecret);
			case "HS384":
				return Algorithm.HMAC384(hsSecret);
			case "HS256":
			default:
				return Algorithm.HMAC256(hsSecret);
			}
		} catch (IllegalArgumentException ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Initialize JWT fail.", ex);
			}
			throw new UserInterfaceJwtErrorException(UserInterfaceJwtErrorException.JwtErrors.JWT_INITIALIZE_FAIL);
		}
	}

	/**
	 * 根据指定的RSA算法名称，构造对应的RSA算法对象
	 *
	 * @param algorithm 算法名称
	 * @return RSA算法对象
	 */
	private Algorithm getRsaAlgorithm(String algorithm) {
		if (StringUtils.isBlank(rsaKeystorePath) || StringUtils.isBlank(rsaPassword1)
				|| StringUtils.isBlank(rsaPassword2) || StringUtils.isBlank(rsaAlias)
				|| StringUtils.isBlank(rsaKeystoreType)) {
			if (logger.isErrorEnabled()) {
				logger.error(String.format(
						"Invalid parameter, keystore: %s, keystore type: %s, password1: %s, password2: %s, alias: %s.",
						rsaKeystorePath, rsaKeystoreType, rsaPassword1, rsaPassword2, rsaAlias));
			}
			throw new UserInterfaceSystemErrorException(
					UserInterfaceSystemErrorException.SystemErrors.SYSTEM_ILLEGAL_PARAM);
		}
		try (FileInputStream fis = new FileInputStream(rsaKeystorePath)) {
			KeyStore keyStore = KeyStore.getInstance(rsaKeystoreType);
			keyStore.load(fis, rsaPassword1.toCharArray());
			RSAPrivateCrtKey privateKey = (RSAPrivateCrtKey) keyStore.getKey(rsaAlias, rsaPassword2.toCharArray());
			RSAPublicKeySpec spec = new RSAPublicKeySpec(privateKey.getModulus(), privateKey.getPublicExponent());
			RSAPublicKey publicKey = (RSAPublicKey) KeyFactory.getInstance("RSA").generatePublic(spec);
			switch (algorithm) {
			case "RSA512":
				return Algorithm.RSA512(publicKey, privateKey);
			case "RSA384":
				return Algorithm.RSA384(publicKey, privateKey);
			case "RSA256":
			default:
				return Algorithm.RSA256(publicKey, privateKey);
			}
		} catch (Exception ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Initialize the JWT fail.", ex);
			}
			throw new UserInterfaceJwtErrorException(UserInterfaceJwtErrorException.JwtErrors.JWT_INITIALIZE_FAIL);
		}
	}

	/**
	 * 初始化相关的资源
	 */
	public void init() {
		try {
			algorithm = getAlgorithm(algorithmStr);

			verifier = JWT.require(algorithm).withIssuer(issuer).withSubject(subject).acceptLeeway(1).acceptExpiresAt(1)
					.build();
		} catch (Exception ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Initialize the JWT fail.", ex);
			}
			throw new UserInterfaceJwtErrorException(UserInterfaceJwtErrorException.JwtErrors.JWT_INITIALIZE_FAIL);
		}
	}

	/**
	 * 销毁相关的资源
	 */
	public void destroy() {
		if (verifier != null) {
			verifier = null;
		}
	}

	/**
	 * {@inheritDoc}
	 *
	 * @see JwtService#signToken(String)
	 */
	@Override
	public String signToken(String accountCode) {
		return signToken(accountCode, null);
	}

	/**
	 * {@inheritDoc}
	 *
	 * @see JwtService#signToken(String, String)
	 */
	@Override
	public String signToken(String accountCode, String expiredTimePeriod) {
		Map<String, Object> claims = new HashMap<>(1);
		claims.put("accountCode", accountCode);
		return signToken(claims, expiredTimePeriod);
	}

	/**
	 * {@inheritDoc}
	 *
	 * @see JwtService#signToken(Map)
	 */
	@Override
	public String signToken(Map<String, Object> claims) {
		return signToken(claims, null);
	}

	/**
	 * {@inheritDoc}
	 *
	 * @see JwtService#signToken(Map, String)
	 */
	@SuppressWarnings("unchecked")
	@Override
	public String signToken(Map<String, Object> claims, String expiredTimePeriod) {
		Date expiredDate = new Date(
				System.currentTimeMillis() + TypeUtils.string2TimePeriod(expiredTimePeriod, 100 * 12 * TypeUtils.MON));
		JWTCreator.Builder builder = JWT.create().withIssuer(issuer).withSubject(subject).withExpiresAt(expiredDate);

		if (builder == null || algorithm == null) {
			throw new UserInterfaceJwtErrorException(UserInterfaceJwtErrorException.JwtErrors.JWT_NOT_INITIALIZE);
		}
		try {
			if (claims != null && !claims.isEmpty()) {
				claims.forEach((k, v) -> {
					if (v instanceof Boolean) {
						builder.withClaim(k, (Boolean) v);
					} else if (v instanceof Integer) {
						builder.withClaim(k, (Integer) v);
					} else if (v instanceof Long) {
						builder.withClaim(k, (Long) v);
					} else if (v instanceof Date) {
						builder.withClaim(k, (Date) v);
					} else if (v instanceof String) {
						builder.withClaim(k, (String) v);
					} else if (v instanceof Double) {
						builder.withClaim(k, (Double) v);
					} else if (v instanceof List) {
						String[] value = ((List<String>) v).toArray(new String[0]);
						builder.withArrayClaim(k, value);
					} else {
						// unsupported type, transform to string
						builder.withClaim(k, v.toString());
					}
				});
			}
			String token = builder.sign(algorithm);
			if (logger.isDebugEnabled()) {
				logger.debug(String.format("Sign the token[%s] successfully.", token));
			}
			return token;
		} catch (Exception ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Sign the token fail.", ex);
			}
			throw new UserInterfaceJwtErrorException(UserInterfaceJwtErrorException.JwtErrors.JWT_SIGN_FAIL);
		}
	}

	/**
	 * {@inheritDoc}
	 *
	 * @see JwtService#verifyToken(String)
	 */
	@Override
	public JwtVerifyResult verifyToken(String token) {
		return verifyToken(token, null);
	}

	/**
	 * {@inheritDoc}
	 *
	 * @see JwtService#verifyToken(String, Predicate)
	 */
	@Override
	public JwtVerifyResult verifyToken(String token, Predicate<Map<String, Claim>> predicate) {
		if (StringUtils.isBlank(token)) {
			throw new UserInterfaceJwtErrorException(UserInterfaceJwtErrorException.JwtErrors.BLANK_TOKEN);
		}
		if (verifier == null) {
			throw new UserInterfaceJwtErrorException(UserInterfaceJwtErrorException.JwtErrors.JWT_NOT_INITIALIZE);
		}
		try {
			DecodedJWT decodedJWT = verifier.verify(token);
			if (predicate == null) {
				return new JwtVerifyResult(decodedJWT);
			} else {
				return predicate.test(decodedJWT.getClaims()) ? new JwtVerifyResult(decodedJWT) : new JwtVerifyResult();
			}
		} catch (Exception ex) {
			if (logger.isErrorEnabled()) {
				logger.error(String.format("Verify the token[%s] fail.", token), ex);
			}
			return new JwtVerifyResult();
		}
	}
}
